#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ComfyUI Linux 安装脚本
基于bash脚本转换的Python版本
"""

import os
import sys
import subprocess
import shutil
import platform
from pathlib import Path
import time

# 导入环境配置 - 使用新的全局变量
from .env_setup import (
    PROJECT_ROOT, COMFYUI_DIR, PYTHON_DIR, PYTHON_CMD,
    get_python_cmd, get_project_root, get_comfyui_dir
)

class ComfyUILinuxInstaller:
    """ComfyUI Linux安装器类"""
    
    # 插件排序
    order = 1
    
    def __init__(self):
        """初始化安装器"""
        # 使用env_setup.py的全局变量
        self.workspace_path = PROJECT_ROOT  # 项目根目录
        self.comfyui_path = COMFYUI_DIR     # ComfyUI目录
        self.config_path = self.workspace_path / "Config"
        self.manager_path = self.config_path / "custom_nodes" / "ComfyUI-Manager"
        self.python_dir = PYTHON_DIR        # Python目录
        self.python_exe = self.python_dir / "bin" / "python3"  # Linux版本使用bin/python3
        
        # 颜色定义
        self.RED = "\033[31m"
        self.GREEN = "\033[32m"
        self.YELLOW = "\033[33m"
        self.BLUE = "\033[34m"
        self.CYAN = "\033[36m"
        self.PURPLE = "\033[35m"
        self.RESET = "\033[0m"
    
    def start(self) -> bool:
        """启动安装器"""
        print(f"\n{self.PURPLE}■■■ ComfyUI Linux 安装器 v1.0 ■■■{self.RESET}")
        print(f"工作目录: {self.CYAN}{self.workspace_path.absolute()}{self.RESET}")
        
        # 执行安装流程
        return self.install_comfyui()
    
    @staticmethod
    def Launch(self):
        """静态启动方法"""
        installer = ComfyUILinuxInstaller()
        return installer.start()
    
    def run_command(self, cmd, shell=True, check=True, capture_output=True, text=True, show_output=True):
        """运行命令并返回结果"""
        try:
            if show_output:
                # 实时显示输出
                print(f"{self.CYAN}执行命令: {cmd}{self.RESET}")
                process = subprocess.Popen(cmd, shell=shell, stdout=subprocess.PIPE, 
                                         stderr=subprocess.STDOUT, text=text, bufsize=1, universal_newlines=True)
                
                # 实时读取输出
                output_lines = []
                while True:
                    if process.stdout is None:
                        break
                    output = process.stdout.readline()
                    if output == '' and process.poll() is not None:
                        break
                    if output:
                        output = output.strip()
                        if output:
                            print(f"  {output}")
                            output_lines.append(output)
                
                # 等待进程完成
                return_code = process.poll()
                if return_code is None:
                    return_code = -1
                result = subprocess.CompletedProcess(cmd, return_code, '\n'.join(output_lines), '')
                return result
            else:
                # 静默模式
                result = subprocess.run(cmd, shell=shell, check=check, 
                                      capture_output=capture_output, text=text)
                return result
        except subprocess.CalledProcessError as e:
            if check:
                print(f"{self.RED}[错误] 命令执行失败: {cmd}{self.RESET}")
                print(f"错误信息: {e.stderr}")
                return None
            else:
                # 当check=False时，返回异常对象作为结果
                return e
        except Exception as e:
            print(f"{self.RED}[错误] 命令执行异常: {cmd}{self.RESET}")
            print(f"异常信息: {e}")
            return None

    def create_project_structure(self):
        """创建AI项目结构"""
        print(f"{self.CYAN}📁 创建 AI 项目结构...{self.RESET}")
        
        # 创建目录
        self.comfyui_path.mkdir(parents=True, exist_ok=True)
        self.config_path.mkdir(parents=True, exist_ok=True)
        
        print(f"{self.GREEN}✅ 项目结构创建完成{self.RESET}")

    def clone_comfyui(self):
        """克隆ComfyUI"""
        print(f"{self.CYAN}🌐 克隆 ComfyUI到：{self.comfyui_path}...{self.RESET}")
        
        # 检查ComfyUI目录是否为空或不存在
        if not self.comfyui_path.exists() or not any(self.comfyui_path.iterdir()):
            print(f"{self.YELLOW}ComfyUI目录为空或不存在，开始克隆...{self.RESET}")
            
            # 如果目录存在但为空，删除它
            if self.comfyui_path.exists():
                shutil.rmtree(self.comfyui_path)
            
            # 尝试克隆，最多重试3次
            for attempt in range(1, 4):
                print(f"{self.CYAN}尝试第 {attempt} 次克隆...{self.RESET}")
                
                # 使用原始GitHub地址
                result = self.run_command(
                    f'git clone --jobs=8 https://github.com/comfyanonymous/ComfyUI.git "{self.comfyui_path}"', 
                    check=False
                )
                
                if result and result.returncode == 0:
                    # 检查是否真的克隆成功了
                    if self.comfyui_path.exists() and any(self.comfyui_path.iterdir()):
                        print(f"{self.GREEN}✅ ComfyUI 克隆成功{self.RESET}")
                        return True
                    else:
                        print(f"{self.RED}❌ ComfyUI 克隆结果为空{self.RESET}")
                else:
                    print(f"{self.RED}❌ ComfyUI 克隆失败 (尝试 {attempt}/3){self.RESET}")
                
                # 如果不是最后一次尝试，等待后重试
                if attempt < 3:
                    print(f"{self.YELLOW}等待 5 秒后重试...{self.RESET}")
                    time.sleep(5)
            
            print(f"{self.RED}❌ ComfyUI 克隆完全失败{self.RESET}")
            return False
        else:
            print(f"{self.GREEN}✅ ComfyUI 已存在且不为空{self.RESET}")
            return True

    def create_virtual_environment(self):
        """创建并激活虚拟环境"""
        print(f"{self.CYAN}🐍 创建并激活虚拟环境...{self.RESET}")
        
        if not self.python_dir.exists():
            result = self.run_command(f'python3 -m venv "{self.python_dir}"', check=False)
            if result and result.returncode == 0:
                print(f"{self.GREEN}✅ 虚拟环境创建成功{self.RESET}")
            else:
                print(f"{self.RED}❌ 虚拟环境创建失败{self.RESET}")
                return False
        else:
            print(f"{self.GREEN}✅ 虚拟环境已存在{self.RESET}")
        
        return True

    def setup_pip_mirrors(self):
        """设置pip镜像"""
        print(f"{self.CYAN}📦 升级pip...{self.RESET}")
        
        # 强制使用默认PyPI源升级pip
        result = self.run_command(
            f'"{self.python_exe}" -m pip install --upgrade pip --index-url https://pypi.org/simple/', 
            check=False
        )
        if result and result.returncode == 0:
            print(f"{self.GREEN}✅ pip 升级成功{self.RESET}")
        else:
            print(f"{self.YELLOW}⚠️ pip 升级失败，继续安装{self.RESET}")
        
        return True

    def clone_comfyui_manager(self):
        """克隆ComfyUI-Manager"""
        print(f"{self.CYAN}🧩 克隆 ComfyUI-Manager 到 custom_nodes...{self.RESET}")
        
        # 确保custom_nodes目录存在
        custom_nodes_dir = self.config_path / "custom_nodes"
        custom_nodes_dir.mkdir(parents=True, exist_ok=True)
        
        if not self.manager_path.exists():
            print(f"{self.CYAN}开始克隆 ComfyUI-Manager...{self.RESET}")
            result = self.run_command(
                f'git clone --jobs=8 https://github.com/Comfy-Org/ComfyUI-Manager.git "{self.manager_path}"', 
                check=False
            )
            if result and result.returncode == 0:
                # 检查是否真的克隆成功了
                if self.manager_path.exists() and any(self.manager_path.iterdir()):
                    print(f"{self.GREEN}✅ ComfyUI-Manager 克隆成功{self.RESET}")
                    return True
                else:
                    print(f"{self.RED}❌ ComfyUI-Manager 克隆结果为空{self.RESET}")
                    return False
            else:
                print(f"{self.RED}❌ ComfyUI-Manager 克隆失败{self.RESET}")
                return False
        else:
            print(f"{self.GREEN}✅ ComfyUI-Manager 已存在{self.RESET}")
            return True

    def create_extra_model_paths(self):
        """创建extra_model_paths.yaml配置文件"""
        print(f"{self.CYAN}📄 自动生成 extra_model_paths.yaml 配置文件...{self.RESET}")
        
        config_content = """custom:
    base_path: "./../Config"
    custom_nodes: "custom_nodes"
    checkpoints: "models/checkpoints"
    clip: "models/clip"
    clip_vision: "models/clip_vision"
    configs: "models/configs"
    controlnet: "models/controlnet"
    diffusers: "models/diffusers"
    diffusion_models: "models/diffusion_models"
    embeddings: "models/embeddings"
    gligen: "models/gligen"
    hypernetworks: "models/hypernetworks"
    loras: "models/loras"
    photomaker: "models/photomaker"
    style_models: "models/style_models"
    unet: "models/unet"
    upscale_models: "models/upscale_models"
    vae: "models/vae"
    vae_approx: "models/vae_approx"
"""
        
        config_file = self.comfyui_path / "extra_model_paths.yaml"
        with open(config_file, 'w', encoding='utf-8') as f:
            f.write(config_content)
        
        print(f"{self.GREEN}✅ extra_model_paths.yaml 配置文件生成完成{self.RESET}")

    def create_model_directories(self):
        """创建模型目录"""
        print(f"{self.CYAN}📁 创建模型目录...{self.RESET}")
        
        model_dir = self.config_path / "models"
        model_dirs = [
            "checkpoints", "clip", "clip_vision", "configs", "controlnet",
            "diffusers", "diffusion_models", "embeddings", "gligen", 
            "hypernetworks", "loras", "photomaker", "style_models", 
            "unet", "upscale_models", "vae", "vae_approx"
        ]
        
        for dir_name in model_dirs:
            dir_path = model_dir / dir_name
            dir_path.mkdir(parents=True, exist_ok=True)
        
        print(f"{self.GREEN}✅ 模型目录创建完成{self.RESET}")

    def create_launcher_script(self):
        """创建启动脚本"""
        print(f"{self.CYAN}📄 创建启动脚本 run_comfyui.sh...{self.RESET}")
        
        launcher_content = f"""#!/bin/bash

# 设置端口（默认设置为 8188，你可以根据需要修改）
PORT=${{1:-8188}}

# 设置Python环境
python_path="{self.python_exe}"

# 查找并杀死占用指定端口（$PORT） 的进程
PID=$(lsof -t -i:$PORT)

if [ ! -z "$PID" ]; then
    echo "发现占用端口 $PORT 的进程（PID: $PID），正在终止进程..."
    kill -9 $PID
    echo "进程已终止"
else
    echo "端口 $PORT 没有被占用"
fi

echo 设置huggingface镜像
export HF_ENDPOINT=https://hf-mirror.com

echo 设置github镜像
# git config --global url."https://gitclone.com/".insteadOf https://

echo 启动后 可以按住ctrl键 点击此链接打开ComfyUI https://bnfviy--$PORT.ap-singapore.cloudstudio.work
# 启动 ComfyUI，指定端口，使用虚拟环境中的python
$python_path "{self.comfyui_path}/main.py" --port "$PORT"
"""
        
        launcher_file = self.workspace_path / "run_comfyui.sh"
        with open(launcher_file, 'w', encoding='utf-8') as f:
            f.write(launcher_content)
        
        # 赋予执行权限
        os.chmod(launcher_file, 0o755)
        
        print(f"{self.GREEN}✅ 启动脚本创建完成{self.RESET}")

    def merge_requirements(self):
        """合并requirements.txt文件"""
        requirements = []
        
        # 读取ComfyUI的requirements.txt
        comfyui_req = self.comfyui_path / "requirements.txt"
        if comfyui_req.exists():
            with open(comfyui_req, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if line and not line.startswith('#'):
                        requirements.append(line)
        
        # 读取Manager的requirements.txt
        manager_req = self.manager_path / "requirements.txt"
        if manager_req.exists():
            with open(manager_req, 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if line and not line.startswith('#'):
                        requirements.append(line)
        
        return requirements

    def install_packages(self, packages, extra_args=""):
        """安装Python包"""
        if packages is None:
            packages = []
        
        failed_packages = []
        
        for i, package in enumerate(packages, 1):
            print(f"\n{self.CYAN}[{i}/{len(packages)}] 安装：{package}{self.RESET}")
            cmd = f'"{self.python_exe}" -m pip install "{package}" {extra_args}'
            result = self.run_command(cmd, check=False, show_output=True)
            
            if result is None or result.returncode != 0:
                print(f"{self.YELLOW}⚠️ 安装失败，尝试安装最新版本...{self.RESET}")
                package_name = package.split('==')[0].split('>=')[0].split('<=')[0].split('~=')[0]
                retry_cmd = f'"{self.python_exe}" -m pip install "{package_name}" --upgrade {extra_args}'
                retry_result = self.run_command(retry_cmd, check=False, show_output=True)
                
                if retry_result is None or retry_result.returncode != 0:
                    print(f"{self.RED}❌ 重试安装也失败{self.RESET}")
                    failed_packages.append(package)
                else:
                    print(f"{self.GREEN}✅ 重试安装成功{self.RESET}")
            else:
                print(f"{self.GREEN}✅ 安装成功{self.RESET}")
        
        # 报告失败的包
        if failed_packages:
            print(f"\n{self.RED}❌ 以下包安装失败：{self.RESET}")
            for pkg in failed_packages:
                print(f"  - {pkg}")
            print(f"\n{self.YELLOW}💡 建议手动安装这些包或检查网络连接{self.RESET}")
        
        return len(failed_packages) == 0

    def detect_cuda(self):
        """检测CUDA版本"""
        print(f"{self.CYAN}🧠 检测是否有独显...{self.RESET}")
        try:
            result = self.run_command("nvidia-smi", check=False, show_output=False)
            if result is not None and result.returncode == 0:
                print(f"{self.GREEN}✅ 检测到 NVIDIA 显卡{self.RESET}")
                
                # 解析CUDA版本
                if result.stdout:
                    for line in result.stdout.split('\n'):
                        if "CUDA Version" in line:
                            version = line.split("CUDA Version:")[1].strip().split()[0]
                            print(f"{self.GREEN}✅ 当前 CUDA 版本: {version}{self.RESET}")
                            
                            # 解析版本号
                            major, minor = version.split('.')[:2]
                            cuda_ver_num = int(major + minor)
                            
                            # 选择PyTorch版本
                            if cuda_ver_num >= 126:
                                return "2.6.0", "https://mirrors.aliyun.com/pytorch-wheels/cu126/"
                            elif cuda_ver_num >= 124:
                                return "2.6.0", "https://mirrors.aliyun.com/pytorch-wheels/cu124/"
                            elif cuda_ver_num >= 121:
                                return "2.5.1", "https://mirrors.aliyun.com/pytorch-wheels/121/"
                            elif cuda_ver_num >= 118:
                                return "2.6.0", "https://mirrors.aliyun.com/pytorch-wheels/cu118/"
                            else:
                                return "cpu", "https://pypi.tuna.tsinghua.edu.cn/simple/"
                
                # 如果没有找到CUDA版本信息，使用默认版本
                return "2.6.0", "https://mirrors.aliyun.com/pytorch-wheels/cu118/"
            else:
                print(f"{self.YELLOW}🧊 未检测到独显，使用 CPU 模式安装 PyTorch...{self.RESET}")
                return "cpu", "https://pypi.tuna.tsinghua.edu.cn/simple/"
        except Exception as e:
            print(f"{self.YELLOW}🧊 检测独显时出错: {e}，使用 CPU 模式安装 PyTorch...{self.RESET}")
            return "cpu", "https://pypi.tuna.tsinghua.edu.cn/simple/"

    def install_comfyui(self) -> bool:
        """执行ComfyUI安装流程"""
        print(f"{self.BLUE}{'=' * 60}{self.RESET}")
        print(f"{self.PURPLE}ComfyUI Linux 安装脚本{self.RESET}")
        print(f"{self.BLUE}{'=' * 60}{self.RESET}")
        
        # 创建项目结构
        self.create_project_structure()
        
        # 克隆ComfyUI
        if not self.clone_comfyui():
            return False
        
        # 创建虚拟环境
        if not self.create_virtual_environment():
            return False
        
        # 设置pip镜像（升级pip并使用默认PyPI源）
        self.setup_pip_mirrors()
        
        # 安装核心依赖（使用默认PyPI源）
        print(f"{self.CYAN}📦 安装核心依赖...{self.RESET}")
        self.run_command(
            f'"{self.python_exe}" -m pip install diffusers transformers --index-url https://pypi.org/simple/', 
            check=False, 
            show_output=True
        )
        
        # 检测CUDA并安装PyTorch
        torch_ver, torch_url = self.detect_cuda()
        print(f"{self.GREEN}✅ 自动选择的 PyTorch 安装版本: {torch_ver}{self.RESET}")
        print(f"{self.CYAN}🌐 安装源: {torch_url}{self.RESET}")
        
        if torch_ver == "cpu":
            print(f"{self.CYAN}📦 安装CPU版本PyTorch...{self.RESET}")
            self.run_command(
                f'"{self.python_exe}" -m pip install torch==2.6.0 torchvision torchaudio --index-url https://pypi.org/simple/ --retries 5', 
                check=False, 
                show_output=True
            )
        else:
            print(f"{self.CYAN}📦 安装GPU版本PyTorch...{self.RESET}")
            self.run_command(
                f'"{self.python_exe}" -m pip install torch=={torch_ver} torchvision torchaudio -f {torch_url} --retries 5', 
                check=False, 
                show_output=True
            )
        
        # 克隆ComfyUI-Manager
        if not self.clone_comfyui_manager():
            return False
        
        # 合并并安装requirements.txt
        print(f"{self.CYAN}🔄 合并安装依赖...{self.RESET}")
        requirements = self.merge_requirements()
        if requirements is None:
            print(f"{self.YELLOW}⚠️ 未找到requirements.txt文件，跳过依赖安装{self.RESET}")
            requirements = []
        else:
            print(f"{self.CYAN}📋 找到 {len(requirements)} 个依赖包{self.RESET}")
        
        # 安装requirements时使用默认PyPI源
        self.install_packages(requirements, "--index-url https://pypi.org/simple/")
        
        # 创建extra_model_paths.yaml
        self.create_extra_model_paths()
        
        # 创建模型目录
        self.create_model_directories()
        
        # 创建启动脚本
        self.create_launcher_script()
        
        # 显示用户信息
        try:
            user = os.getenv('USER', 'unknown')
            print(f"{self.CYAN}当前用户: {user}{self.RESET}")
        except:
            pass
        
        # 安装完成
        print(f"\n{self.BLUE}{'=' * 60}{self.RESET}")
        print(f"{self.GREEN}✅ 安装完成！您可以使用 './run_comfyui.sh' 启动 ComfyUI。{self.RESET}")
        print(f"{self.BLUE}{'=' * 60}{self.RESET}")
        
        return True

def main():
    """主函数"""
    installer = ComfyUILinuxInstaller()
    return installer.start()

if __name__ == "__main__":
    sys.exit(0 if main() else 1) 